import pandas as pd
import numpy as np
import torch.nn.functional as F
import cv2
import os
import torch
from data_aug import resize
from process_data import dicom2array
from models.model import model_unet,model_cpn50,model_simple_pose_res18,model_simple_pose_res101,model_resnet_unet,model_DLinknet34
from dataset import heatmap_to_anno,convert_dict,conver_csv,multi_heatmap_to_anno,heatmap_to_anno_with_category

def compute_loss(pred,label,original_size,resize_size,distance = 1):
    ## 用来记录距离小于指定数值的正样本个数
    positive = 0
    ## 用来记录l2 loss
    l2_loss = []
    assert len(pred) == len(label),'the len of pred and label shoule be same'
    ## 输入全部的pred和label，用list中包含tuple的形式，其中像素间距为0.139mm
    for index in range(len(pred)):
        single_pred = pred[index]
        single_label = label[index]
        ## 遇到缺失的label不进行计算
        if 0 in single_label:
            continue
        ## 计算两个点之间的像素距离。注意根据原始尺寸和resize后的尺寸对坐标值进行转换。
        loss = (((single_pred[0]-single_label[0])/resize_size[0]*original_size[0])**2 +
                ((single_pred[1]-single_label[1])/resize_size[1]*original_size[1])** 2)**0.5
        ## 将计算出的像素距离转换成物理距离
        distance_loss = float(loss)*0.139
        # print('distance loss',distance_loss)
        l2_loss.append(distance_loss)
        if distance_loss <= distance:
            positive +=1
    return positive,l2_loss

def test_aug_fun(model,input):
    ## 测试增强操作。输入转成tensor的数据。
    input_1 = input
    input_2 = torch.flip(input,dims=[3])
    out_1 = model(input_1)
    out_2 = model(input_2)
    out_2 = torch.flip(out_2,dims=[3])
    out = (out_1 + out_2) / 2
    return out



def predict(model,image_dir,size=(512,512),show_save = None,test_aug = False,model2 = None):
    point_name_list = ['L1', 'L2', 'L3', 'L4', 'L5', 'T12-L1', 'L1-L2', 'L2-L3', 'L3-L4', 'L4-L5', 'L5-S1']
    with torch.no_grad():
        model.eval()
        image_dir = image_dir.replace('\\','/')
        if 'dcm' in image_dir:
            image = dicom2array(image_dir)
            image = np.stack([image,image,image],-1)
        else:
            image = cv2.imread(image_dir)
        original_size = (image.shape[:-1])[::-1]
        original_image = image.copy()
        # resize the data
        image = cv2.resize(image,size[::-1])
        input = torch.from_numpy(np.transpose(image, axes=(2, 0, 1))).float().cuda()
        input = input[np.newaxis, ...]
        if test_aug:
            pred = test_aug_fun(model, input)
        else:
            pred = model(input)
        # pred = test_aug(model,input)
        ## 注意pred[0]去掉batch维度
        ## 返回的坐标为xy的形式。
        pred = heatmap_to_anno(pred[0])
        # pred = multi_heatmap_to_anno(pred[0])
        pred = [[int(p[0] / size[0] * original_size[0]), int(p[1] / size[1] * original_size[1])] for p in pred]
        pred_dict = dict(zip(point_name_list, pred))

        if show_save:
            for index in range(len(pred)):
                single_pred = pred[index]
                ##  注意此时坐标已经经过转换处理了
                cv2.circle(original_image, (single_pred[0],single_pred[1]), 1, (128, 255, 0), thickness=4)
            # cv2.imshow('image',image)
            # cv2.waitKey(0)
            cv2.imwrite(os.path.join(show_save,image_dir.split('/')[-1]), original_image)

        return pred_dict




def val_fun(model,anno_list,image_size,image_save_path = '../data/image',distance = 1.,
            show_pred = False,test_aug = False,model2 = None):
    all_l2_loss = []
    all_positive = 0
    num_image = len(anno_list)
    with torch.no_grad():
        model.eval()
        for index in range(num_image):
            try:
                image_path = anno_list[index]['image_path']
                image_path = os.path.join(image_save_path, image_path.split('//')[1] + '_' +
                                          image_path.split('//')[2].replace('dcm', 'jpg'))
                image = cv2.imread(image_path)
                original_size = image.shape[:-1]
                ## 读取label
                label = anno_list[index]
                label = convert_dict(label)
                anno = [p[0] for p in label]
                category = [p[1] for p in label]
                spacing = [p[2] for p in label]
                category_1 = [int('v2' in p) for p in category[:5]]
                category_2 = [int(p[1]) - 1 for p in category[5:]]
                # resize the data
                image, anno = resize(image, anno, image_size)
                input = torch.from_numpy(np.transpose(image,axes=(2,0,1))).float().cuda()
                input = input[np.newaxis,...]

                if test_aug:
                    pred = test_aug_fun(model,input)
                else:
                    pred = model(input)
                # cv2.imshow('heatmap1', pred[0][0].cpu().numpy().astype(np.uint8))
                # cv2.waitKey()
                # cv2.imshow('heatmap2', pred[0][1].cpu().numpy().astype(np.uint8))
                # cv2.waitKey()
                ## 注意pred[0]去掉batch维度
                pred_anno = pred
                pred_anno = heatmap_to_anno(pred_anno[0])
                # pred = multi_heatmap_to_anno(pred[0])
                # print('pred anno',pred)
                # print('label anno',label)
                # print('----------------------------')

                if show_pred:
                    for index in range(len(pred_anno)):
                        single_pred = pred_anno[index]
                        single_anno = anno[index]
                        single_anno = (int(single_anno[0]),int(single_anno[1]))
                        if 0 in single_anno :
                            continue
                        ##  注意此时坐标已经经过转换处理了
                        cv2.circle(image, single_anno[:],1,(255,0,255),thickness=5)
                        cv2.circle(image, single_pred[:], 1, (255, 255, 0), thickness=2)
                    # cv2.imshow('image',image)
                    # cv2.waitKey(0)
                    cv2.imwrite(image_path.replace(image_save_path,'../data/show/pred',),image)
                positive,l2_loss = compute_loss(pred_anno,anno,original_size=original_size,resize_size=image_size,distance=distance)
                all_positive += positive
                all_l2_loss = all_l2_loss+l2_loss
            except Exception as e:
                # print(e)
                continue
    ## 由于部分label存在缺失，不进行loss的计算。根据all_l2_loss的长度来作为实际参与计算的点的数量。

    return all_positive/len(all_l2_loss),np.mean(all_l2_loss),np.median(all_l2_loss)



if __name__ == '__main__':

    # Image_Size = (512,512)
    Image_Size = (768,768)
    # Image_Size = (1024, 1024)
    Double_Cate = True
    Test_aug = True
    Num_Class = 2
    val_csv_dir = '../data/csv_label/train51_anno.csv'
    # trained_model_dir = '../model_save/DlinkNet34/best_.pth'
    trained_model_dir = '../model_save/new/DlinkNet34_768/best_.pth'
    trained_model_dir_2 = '../model_save/DlinkNet34_768/best_.pth'
    # model = model_cpn50(image_size = Image_Size)
    # model = model_simple_pose_res101()
    # model = model_resnet_unet(layer='resnet34',pre_train=False,num_class=Num_Class)
    model = model_DLinknet34(num_class=Num_Class,pre_trained=False)
    model_state = torch.load(trained_model_dir)
    model_state = {k:v for k,v in model_state.items() if 'classify' not in k}
    model.load_state_dict(model_state)

    # model_2 = model_DLinknet34(num_class=Num_Class, pre_trained=False, double=Double_Cate)
    # model_2_state = torch.load(trained_model_dir_2)
    # model_2_state = {k:v for k,v in model_2_state.items() if 'classify' not in k }
    # model_2.load_state_dict(model_2_state)

    recall, l2_mean, l2_medium,  = val_fun(model, val_csv_dir, Image_Size, distance=1,
                                                                 show_pred=True,test_aug=False,model2=None)
    print('recall', recall, 'l2 mean', l2_mean, 'l2 medium', l2_medium,)
